from flask import Flask, render_template, request, jsonify, send_from_directory
import os
from werkzeug.utils import secure_filename
import traceback
import time
import json
from back import DocumentProcessor

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB限制
app.secret_key = 'your-secret-key-here'

os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

processor = DocumentProcessor()

def extract_structured_data(raw_data):
    """从原始API响应中提取结构化数据"""
    if isinstance(raw_data, dict) and 'choices' in raw_data:
        try:
            return json.loads(raw_data['choices'][0]['message']['content'])
        except (json.JSONDecodeError, KeyError, AttributeError):
            return []
    elif isinstance(raw_data, list):
        return raw_data
    return []

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_files():
    if 'requirement' not in request.files or 'test' not in request.files:
        return jsonify({'success': False, 'error': '请上传需求文档和测试报告'}), 400
    
    req_file = request.files['requirement']
    test_file = request.files['test']
    
    if req_file.filename == '' or test_file.filename == '':
        return jsonify({'success': False, 'error': '未选择文件'}), 400
    
    try:
        timestamp = str(int(time.time()))
        req_filename = f"req_{timestamp}_{secure_filename(req_file.filename)}"
        test_filename = f"test_{timestamp}_{secure_filename(test_file.filename)}"
        
        req_path = os.path.join(app.config['UPLOAD_FOLDER'], req_filename)
        test_path = os.path.join(app.config['UPLOAD_FOLDER'], test_filename)
        
        req_file.save(req_path)
        test_file.save(test_path)
        
        return jsonify({
            'success': True,
            'req_filename': req_filename,
            'test_filename': test_filename
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': f'文件保存失败: {str(e)}'
        }), 500

@app.route('/process', methods=['POST'])
def process_documents():
    data = request.get_json()
    if not data or 'req_filename' not in data or 'test_filename' not in data:
        return jsonify({'success': False, 'error': '无效的请求数据'}), 400
    
    req_path = os.path.join(app.config['UPLOAD_FOLDER'], data['req_filename'])
    test_path = os.path.join(app.config['UPLOAD_FOLDER'], data['test_filename'])
    
    try:
        raw_results = processor.full_processing_pipeline(req_path, test_path)
        
        # 提取并清理数据
        use_cases = extract_structured_data(raw_results.get('use_cases', []))
        test_results = extract_structured_data(raw_results.get('test_results', []))
        
        # 标准化数据结构
        standardized_use_cases = []
        for item in use_cases:
            standardized_use_cases.append({
                'id': item.get('ID', ''),
                'name': item.get('Name', '未命名需求'),
                'type': item.get('Type', '未指定类型'),
                'description': item.get('Description', '无描述'),
                'priority': item.get('Priority', '未指定'),
                'source': item.get('Source', '')
            })
        
        standardized_test_results = []
        for item in test_results:
            standardized_test_results.append({
                'id': item.get('ID', ''),
                'name': item.get('Name', '未命名测试'),
                'requirement_id': item.get('RequirementID', '未关联'),
                'objective': item.get('TestObjective', ''),
                'method': item.get('TestMethod', ''),
                'result': item.get('TestResult', '无结果'),
                'status': item.get('Status', '未知'),
                'evidence': item.get('Evidence', ''),
                'defect_id': item.get('DefectID', '')
            })
        
        return jsonify({
            'success': True,
            'data': {
                'use_cases': standardized_use_cases,
                'test_results': standardized_test_results
            }
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e),
            'traceback': traceback.format_exc()
        }), 500

@app.after_request
def after_request(response):
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
    response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
    return response

@app.route('/download-report', methods=['POST'])
def download_report():
    data = request.get_json()
    if not data or 'req_filename' not in data or 'test_filename' not in data:
        return jsonify({'success': False, 'error': '无效的请求数据'}), 400
    
    try:
        # 从请求文件名中提取文档名称
        req_filename = data['req_filename']
        doc_name = Path(req_filename).stem.replace("req_", "").split('_', 1)[-1]
        
        # 构建报告文件路径
        output_dir = Path("processed_docs") / doc_name
        report_path = output_dir / "extracted_results.txt"
        
        # 验证报告文件是否存在
        if not report_path.exists():
            raise FileNotFoundError(f"报告文件不存在于: {report_path}")
            
        return send_from_directory(
            directory=output_dir,
            path="extracted_results.txt",
            as_attachment=True,
            download_name=f"{doc_name}_解析报告.txt"
        )
        
    except FileNotFoundError as e:
        return jsonify({
            'success': False,
            'error': f"找不到报告文件，请重新处理文档。原因为: {str(e)}"
        }), 404
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e),
            'traceback': traceback.format_exc()
        }), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)